#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
绘制参考基因组覆盖率图（按固定窗口分段统计 reads hit 数）

功能：
1. 读取 litesam4ref 解析各物种的 reads 位置
2. 读取 ref_length 文件获得参考序列长度
3. 生成分段窗口覆盖率统计
4. 输出每个物种的覆盖率直方图 PNG

作者: Bing Liang
日期: 2025-11-19
"""

import argparse
import os
from collections import defaultdict

import matplotlib.pyplot as plt
import pandas as pd
from pathlib import Path


# ----------------------------------------------------------------------
# 数据解析函数
# ----------------------------------------------------------------------
def parse_lites_reads(job_id, sample_id) -> dict:
    """
    读取 *.litesam4ref 文件，格式：
        species,12-100,56-100,...
    返回 {species: [position, ...]}
    """
    out = defaultdict(list)

    with open(Path(f"/data/mNGS/runmngs/result/{job_id}/{sample_id}/{sample_id}.litesam4ref"), "r") as fr:
        for line in fr:
            if line := line.strip():
                sp, *rest = line.split(",")
                for item in rest:
                    pos, _ = item.split("-")
                    out[sp].append(int(pos))

    return out


def parse_ref_length() -> dict:
    """
    读取 ref_length_4_coverage.xls 文件（两列：物种\t长度）
    返回 {species: length}
    """
    out = {}
    with open(Path("/data/mNGS/pipeline/script/table/ref_length_4_coverage.xls"), "r") as fr:
        for line in fr:
            if line := line.strip():
                items = line.split("\t")
                out[items[0]] = int(items[1])
    return out


# ----------------------------------------------------------------------
# 覆盖率统计
# ----------------------------------------------------------------------
def compute_windows(ref_len: int, positions: list, num_windows: int = 100) -> pd.DataFrame:
    """
    将基因组按固定窗口切分，并统计每个窗口命中 reads 数
    """
    width = ref_len // num_windows
    out = defaultdict(list)

    for i in range(1, num_windows + 1):
        start = (i - 1) * width + 1
        end = ref_len if i == num_windows else i * width

        # 命中统计
        if i == num_windows:
            count = len([p for p in positions if start <= p <= end])
        else:
            count = len([p for p in positions if start <= p < end])

        out["Window"].append(i)
        out["Start"].append(start)
        out["End"].append(end)
        out["Count"].append(count)

        # 每 10 个窗口设置一次标签：kb 单位
        out["Label"].append((start // 1000) * 10)

    return pd.DataFrame(out)


# ----------------------------------------------------------------------
# 绘图
# ----------------------------------------------------------------------
def plot_coverage(df: pd.DataFrame, species: str, outdir: str, dpi: int = 300):
    """
    绘制 coverage 的直方图
    """
    plt.figure(figsize=(8, 4))
    plt.bar(df["Window"], df["Count"], width=0.825, color="#00AA00")

    # X 轴标签（每 10 个窗口）
    xticks = list(range(10, 91, 10))
    xlabels = df.loc[df["Window"].isin(xticks), "Label"].tolist()
    plt.xticks(xticks, xlabels)

    ax = plt.gca()
    ax.yaxis.get_major_locator().set_params(integer=True)

    # X/Y 轴从 0 开始
    ax.spines['left'].set_position(('data', 0))
    ax.spines['bottom'].set_position(('data', 0))

    plt.xlabel("Genome Position (kb)")
    plt.ylabel("Reads Count")
    plt.title(species)
    plt.margins(x=0.02)

    os.makedirs(outdir, exist_ok=True)
    outfile = os.path.join(outdir, f"{species}.png")

    plt.savefig(outfile, dpi=dpi, bbox_inches="tight")
    plt.close()


# ----------------------------------------------------------------------
# 主流程
# ----------------------------------------------------------------------
def main(args):
    sp_reads = parse_lites_reads(args.job_id, args.sample_id)
    sp_len = parse_ref_length()

    for sp, positions in sp_reads.items():
        if sp not in sp_len:
            print(f"[Warning] {sp} 不在 ref_length 文件中，跳过")
            continue

        ref_len = sp_len[sp]
        df = compute_windows(ref_len, positions, args.windows)

        plot_coverage(df, sp, args.outdir, args.dpi)
        print(f"[OK] {sp} 覆盖图已生成")


# ----------------------------------------------------------------------
# 命令行参数
# ----------------------------------------------------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="基因组窗口覆盖率绘图工具")

    parser.add_argument("--job_id", required=True, help="芯片编号")
    parser.add_argument("--sample_id", required=True, help="样本编号")
    parser.add_argument("--outdir", default="coverage_plots", help="输出目录")
    parser.add_argument("--windows", type=int, default=100, help="窗口数量")
    parser.add_argument("--dpi", type=int, default=300, help="输出图像 dpi")

    main(parser.parse_args())
